#Copyright (c) 2025, Alibaba Cloud and its affiliates;
#Licensed under the Apache License, Version 2.0 (the "License");
#you may not use this file except in compliance with the License.
#You may obtain a copy of the License at

#   http://www.apache.org/licenses/LICENSE-2.0

#Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and
#limitations under the License.


import json
import random
import string
from datetime import datetime

from flask import current_app as app

from services.common import CommonService
from common_libs.helper import response_data
import common_libs.constants as constants
from extensions.exception import FceException
from common_libs.sqlitedb import FCE_DB, RayClusters, RayClustersDeleted
import common_libs.sqlitedb as db 


class ClusterService(CommonService):
    def __init__(self):
        return

    def init(self, req):
        if not req['SDKPath']:
            app.logger.error('RequestId:{}, SDKPath is required'.format(req['RequestId']))
            FceException('InvalidParameters.Empty', 'SDKPath')
    
        if not req['KubeConfPath']:
            app.logger.error('RequestId:{}, KubeConfPath is required'.format(req['RequestId']))
            FceException('InvalidParameters.Empty', 'KubeConfPath')

        cmd = "deploy tool_check "
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            cmd = "deploy tool_install --mode local --pkg_path {}/tools --install_path {} ".format(req['SDKPath'], constants.FCE_TOOLS_INSTALL_PATH)
            ret = self.exec_fceutil_cmd(cmd)
            if ret["Success"] == False:
                FceException('Deploy.ToolInstallFailed', ret)

        cmd = "deploy kube_check --conf_path {} ".format(req['KubeConfPath'])
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            FceException('Deploy.KubeNotFound')
        app.logger.debug('RequestId:{}, kube_check response is {}'.format(req['RequestId'], response))

        check_result = self.deploy_kuberay_operator(req)
        if check_result['Success'] == False:
            FceException('Deploy.OperatorInstallFailed', {check_result['Message']['ErrMsg']})
        
        data = {
            "RequestId": req['RequestId'], 
        }
        return response_data(True, data, "Success")

    def list_ray_cluster(self,req):
        if not req['KubeConfPath']:
            app.logger.error('RequestId:{}, KubeConfPath is required'.format(req['RequestId']))
            FceException('InvalidParameters.Empty', 'KubeConfPath')
        
        query = db.RayClusters.query.filter(RayClusters.namespace == req['Namespace'])
        if len(req['ClusterNames']) > 0:
            query = query.filter(RayClusters.cluster_name.in_(req['ClusterNames']))
        if req['ClusterId']:
            query = query.filter(RayClusters.cluster_id == req['ClusterId'])
        if req['ClusterVersion']:
            query = query.filter(RayClusters.cluster_version == req['ClusterVersion'])

        clusters = query.all()
        if len(clusters) == 0:
            return response_data(True, {"RequestId": req['RequestId'], "TotalCount": 0, "ClusterList": []}, "Success")
        
        app.logger.error('RequestId:{}, db query result is {}'.format(req['RequestId'], clusters))
        for cluster in clusters:
            app.logger.debug('RequestId:{}, cluster_name: {}, cluster_id: {}'.format(req['RequestId'], cluster.cluster_name, cluster.cluster_id))
        cluster_dict = {cluster.cluster_name: {'cluster_id': cluster.cluster_id, 'create_time': cluster.create_time}for cluster in clusters}
        app.logger.info('cluster_dict: {}'.format(cluster_dict))
        #construct the fce command line
        cmd = "rayc cluster_list"
        if len(req['ClusterNames']) > 0:
            cmd = cmd + " --cluster_name " + " ".join(req.ClusterNames)
        if req['ClusterId']:
            cmd = cmd + " --cluster_id " + req['ClusterId']
        if req['Namespace']:
            cmd = cmd + " --namespace " + req['Namespace']
        if req['KubeConfPath']:
            cmd = cmd + " --kube_config " + req['KubeConfPath']
        app.logger.info('RequestId:{}, list ray cluster cmd: {}'.format(req['RequestId'], cmd))

        #Run fce command line and return analysis result
        response = self.exec_fceutil(cmd)
        app.logger.info('response: {}'.format(response))
        for cluster in response['ClusterList']:
            if cluster['ClusterName'] in cluster_dict:
                cluster['ClusterId'] = cluster_dict[cluster['ClusterName']]['cluster_id']
                cluster['CreateTime'] = cluster_dict[cluster['ClusterName']]['create_time']

            else:
                app.logger.warn('RequestId:{}, cluster:{} is not in fce database'.format(req['RequestId'], cluster['ClusterName']))
                response['ClusterList'].remove(cluster)
        total_count = len(response['ClusterList'])
        app.logger.info('response after: {}'.format(response))
        
        data = {
            "RequestId": req['RequestId'], 
            "TotalCount": total_count,
            "ClusterList": response['ClusterList'],
        }
        return response_data(True, data, "Success")

    def gen_cluster_id(self, prefix, table):
        for _ in range(100):
            random_part = ''.join(random.choices(string.ascii_lowercase + string.digits, k=6))
            id = '-'.join([prefix, random_part])
            result = db.FCE_DB.session.query(table).filter(table.cluster_id == id).all()
            if len(result) == 0:
                return id

    def create_ray_cluster_obj(self, req):
        cmd = "rayc cluster_create --kube_config {} --cluster_id {} --cluster_name {} --namespace {} --head_node '{}' --worker_node '{}'".format(req['KubeConfPath'], req['ClusterId'], req['ClusterName'], req['Namespace'], json.dumps(req['HeadNode']), json.dumps(req['WorkerNodeGroups']))
    
        return self.exec_fceutil_cmd(cmd, 30)

    def deploy_kuberay_operator(self, req):
        cmd = "rayc operator_check --kube_config {} --namespace {}".format(req['KubeConfPath'], req['Namespace'])
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            app.logger.warn('RequestId:{}, kube ray operator check failed'.format(req['RequestId']))
            cmd = "rayc operator_create --kube_config {} --namespace {} --image {}".format(req['KubeConfPath'], req['Namespace'], req['OperatorImage'])
            response = self.exec_fceutil_cmd(cmd, 30)
            if response["Success"] == False:
                app.logger.error('RequestId:{},  kube ray operator deploy failed'.format(req['RequestId']))
                #FceException('RayCluster.OperatorInstallFailed')
        return response

    def deploy_cluster_dependency_resources(self, req):
        cmd = "deploy res_obj --conf_path {} --cluster_id {} --namespace {}".format(req['KubeConfPath'], req['ClusterId'], req['Namespace'])
    
        return self.exec_fceutil_cmd(cmd)
      
    def clean_cluster_resources(self, req):
        cmd = "clean res_obj --conf_path {} --cluster_id {} --cluster_name {} --namespace {}".format(req['KubeConfPath'], req['ClusterId'], req['ClusterName'], req['Namespace'])
    
        return self.exec_fceutil_cmd(cmd)

    def delete_ray_cluster_obj(self, req):
        cmd = "rayc cluster_delete --kube_config {} --cluster_id {} --cluster_name {} --namespace {} ".format(req['KubeConfPath'], req['ClusterId'], req['ClusterName'], req['Namespace'])
    
        return self.exec_fceutil_cmd(cmd)

    def create_ray_cluster(self, req):
        #1.parameters validation
        if not req['ClusterName']:
            app.logger.error('RequestId:{}, ClusterName is required'.format(req['RequestId']))
            FceException('InvalidParameters.Empty', 'ClusterName')
    
        if not req['KubeConfPath']:
            app.logger.error('RequestId:{}, KubeConfPath is required'.format(req['RequestId']))
            FceException('InvalidParameters.Empty', 'KubeConfPath')

        result = db.RayClusters.query.filter(RayClusters.cluster_name == req['ClusterName']).first()
        app.logger.error('RequestId:{}, db query result is {}'.format(req['RequestId'], result))

        if result:
            app.logger.error('RequestId:{}, ClusterName {} already exists'.format(req['RequestId'], req['ClusterName']))
            FceException('InvalidParameters.AlreadyExists', 'ClusterName', req['ClusterName'])

        ray_cluster_id = self.gen_cluster_id(req['ClusterName'], RayClusters)
        app.logger.info('RequestId:{}, ray cluster id: {}'.format(req['RequestId'], ray_cluster_id))
        req['ClusterId'] = ray_cluster_id

        '''
        ###To be moved
        cmd = "deploy tool_check "
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            cmd = "deploy tool_install --mode local --pak_path {} --install_path {} ".format(constants.FCE_TOOLS_PKG_PATH, constants.FCE_TOOLS_INSTALL_PATH)
            ret = self.exec_fceutil_cmd(cmd)
            if ret["Success"] == False:
                FceException('InvalidParameters.Empty', 'KubeConfPath')

        cmd = "deploy kube_check --conf_path {} ".format(req['KubeConfPath'])
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            FceException('Deploy.KubeNotFound')
        req['KubeClusterId'] = response['KubeClusterId']
        req['KubeClusterUrl'] = response['KubeClusterUrl']

        check_result = self.deploy_kuberay_operator(req)
        if check_result['Success'] == False:
            FceException('Deploy.OperatorCheckFailed', deploy_result)
        ###
        '''
        cmd = "deploy tool_check "
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            FceException('Deploy.ToolCheckFailed')

        cmd = "deploy kube_check --conf_path {} ".format(req['KubeConfPath'])
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            FceException('Deploy.KubeNotFound')
        app.logger.debug('RequestId:{}, kube_check response is {}'.format(req['RequestId'], response))
        req['KubeClusterId'] = response['Message']['KubeClusterId']
        req['KubeClusterUrl'] = response['Message']['KubeClusterUrl']

        cmd = "rayc cluster_list --kube_config {} --cluster_name {} ".format(req['KubeConfPath'], req['ClusterName'])
        if req['Namespace']:
            cmd = cmd + " --namespace " + req['Namespace']
        response = self.exec_fceutil_cmd(cmd)
        if response['Success'] == True:
            FceException('RayCluster.NameAlreadyExist')
            
        deploy_result = self.deploy_cluster_dependency_resources(req)
        if deploy_result['Success'] == False:
            FceException('RayCluster.ResourcesDeployFailed', deploy_result)

        create_result = self.create_ray_cluster_obj(req)
        if create_result['Success'] == False:
            FceException('RayCluster.ClusterInstallFailed', create_result)
        req['Status'] = create_result['Message']['Status']
        app.logger.debug('RequestId:{}, create ray cluster object success'.format(req['RequestId']))

        try:
            cluster_info = RayClusters(cluster_name=req['ClusterName'], cluster_version=req['ClusterVersion'], cluster_id=req['ClusterId'],
                                   cluster_type='Ray', head_node=json.dumps(req['HeadNode']),
                                   worker_groups=json.dumps(req['WorkerNodeGroups']), kube_config=req['KubeConfPath'], kube_cluster_id=req['KubeClusterId'],
                                   namespace=req['Namespace'], status=req['Status'],
                                   create_time=datetime.now())
            db.insert_db(cluster_info)
            #return response_data(True, {"Message": "RequestId %s: success" % req['RequestId']}, "Success")
        except Exception as e:
            FceException('OperationFailed.Database', e)

        data = {
            "RequestId": req['RequestId'], 
            "ClusterName": req['ClusterName'],
            "ClusterId": req['ClusterId'],
            "KubeClusterId": req['KubeClusterId'],
        }
        return response_data(True, data, "Success")

    def delete_ray_cluster(self, req):
        if not req['ClusterId']:
            FceException('InvalidParameters.Empty', 'ClusterId')
    
        if not req['KubeConfPath']:
            FceException('InvalidParameters.Empty', 'KubeConfPath')

        result = db.RayClusters.query.filter(RayClusters.cluster_id == req['ClusterId']).first()
        if not result:
            FceException('RayCluster.ClusterNotExist', 'ClusterId = ' + req['ClusterId'])
        req['ClusterName'] = result.cluster_name

        cmd = "deploy tool_check "
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            FceException('Deploy.ToolCheckFailed')

        cmd = "deploy kube_check --conf_path {} ".format(req['KubeConfPath'])
        response = self.exec_fceutil_cmd(cmd)
        if response["Success"] == False:
            FceException('Deploy.KubeNotFound')
        #req['KubeClusterId'] = response['KubeClusterId']

        cmd = "rayc cluster_list --kube_config {} --cluster_name {} ".format(req['KubeConfPath'], req['ClusterName'])
        if req['Namespace']:
            cmd = cmd + " --namespace " + req['Namespace']
        
        response = self.exec_fceutil_cmd(cmd)
        if response['Success'] != True:
            FceException('RayCluster.ClusterNotExist', response)

        delete_result = self.delete_ray_cluster_obj(req)
        if delete_result['Success'] == False:
            FceException('RayCluster.DeleteClusterFailed', delete_result)

        clean_result = self.clean_cluster_resources(req)
        if clean_result['Success'] == False:
            FceException('RayCluster.CleanResourcesFailed', clean_result)

        try:
            delete_item = result
            app.logger.info(delete_item)
            db.delete_db(delete_item)

            delete_item_add = RayClustersDeleted(cluster_name=delete_item.cluster_name, cluster_version=delete_item.cluster_version, cluster_id=delete_item.cluster_id, cluster_type=delete_item.cluster_type, head_node=delete_item.head_node, worker_groups=delete_item.worker_groups, kube_config=delete_item.kube_config, kube_cluster_id=delete_item.kube_cluster_id, namespace=delete_item.namespace, status=constants.ClusterStatus.DELETED.value, create_time=delete_item.create_time, delete_time=datetime.now())

            db.insert_db(delete_item_add)
            #return response_data(True, {"Message": "RequestId %s: success" % req['RequestId']}, "Success")
        except Exception as e:
            FceException('OperationFailed.Database', e)

        data = {
            "RequestId": req['RequestId'], 
        }
        return response_data(True, data, "Success")
    
    def get_ray_cluster(self,req):
        if not req['ClusterId']:
            FceException('InvalidParameters.Empty', 'ClusterId')

        if not req['KubeConfPath']:
            app.logger.error('RequestId:{}, KubeConfPath is required'.format(req['RequestId']))
            FceException('InvalidParameters.Empty', 'KubeConfPath')
        
        result = db.RayClusters.query.filter(RayClusters.cluster_id == req['ClusterId']).first()
        if not result:
            FceException('RayCluster.ClusterNotExist', 'ClusterId = ' + req['ClusterId'])
        req['ClusterName'] = result.cluster_name
        req['Namespace'] = result.namespace

        try:
            head_node_info = json.loads(result.head_node)
        except Exception as e:
            FceException('InvalidDataType.FceutilResult', e)

        try:
            worker_groups = json.loads(result.worker_groups)
        except Exception as e:
            FceException('InvalidDataType.FceutilResult', e)

        
        app.logger.error('RequestId:{}, db query result is {}'.format(req['RequestId'], result))
        cmd = "rayc cluster_get --kube_config {} --cluster_name {} ".format(req['KubeConfPath'], req['ClusterName'])
        if req['Namespace']:
            cmd = cmd + " --namespace " + req['Namespace']
        app.logger.info('RequestId:{}, get ray cluster cmd: {}'.format(req['RequestId'], cmd))

        #Run fce command line and return analysis result
        response = self.exec_fceutil(cmd)
        app.logger.info('response: {}'.format(response))
        cluster = response['ClusterInfo']
        cluster['HeadNode'] = head_node_info
        cluster['WorkerGroupNode'] = worker_groups
        cluster['ClusterId'] = result.cluster_id

        data = {
            "RequestId": req['RequestId'], 
            "ClusterInfo": cluster,
        }
        return response_data(True, data, "Success")

clusterService = ClusterService()